from dataLoader import Data
from skfuzzy.cluster import cmeans
from vision import plot_loss, plot_scatter, plot_pdf


class Trainer:
    def __init__(self, args):
        self.args = args

        self.dataset = Data()
        self.file = open(self.args.path + 'result.txt', 'w')
        self.write('结果', 1)

    def write(self, data, next_line=1):
        self.file.write(data)
        for i in range(next_line):
            self.file.write('\n')
        self.file.flush()

    def save_cluster_result(self, label):
        df = self.dataset.df
        df['label'] = label
        df.to_csv(self.args.path + 'cluster_results.csv')

    def start(self):
        data = self.dataset.data

        center, u, u0, d, jm, p, fpc = cmeans(data.T, m=2, c=self.args.c, error=0.00005, maxiter=self.args.maxiter)
        plot_loss(jm, self.args.path)

        label = u.T.argmax(1)
        plot_scatter(data, center, label, self.args.path)
        self.write('聚类中心:', 1)
        self.write(str(center), 2)
        self.write('各个簇的个数:')
        for i in range(self.args.c):
            self.write('类%d: %d' % (i, (label == i).sum()))
            plot_pdf(data[label == i], 'DST', self.args.path, 'DST_类别%d.png' % i)

            plot_pdf(data[label == i], 'PTDTC', self.args.path, 'PTDTC_类别%d.png' % i)

        self.file.close()

        self.save_cluster_result(label)
